"""
数据读取模块
"""
import pandas as pd
import psycopg2
from typing import Dict, List, Any, Optional
import logging
from config import db_config

logger = logging.getLogger(__name__)

class DataReader:
    """数据读取器"""
    
    def __init__(self, connection_name: str = 'default'):
        self.connection_name = connection_name
        self.config = db_config.get_connection_config(connection_name)
        self.connection = None
    
    def connect(self):
        """建立数据库连接"""
        try:
            self.connection = psycopg2.connect(
                host=self.config['host'],
                port=self.config['port'],
                database=self.config['database'],
                user=self.config['username'],
                password=self.config['password']
            )
            logger.info(f"成功连接到数据库: {self.config['database']}")
        except Exception as e:
            logger.error(f"数据库连接失败: {str(e)}")
            raise
    
    def disconnect(self):
        """关闭数据库连接"""
        if self.connection:
            self.connection.close()
            logger.info("数据库连接已关闭")
    
    def read_table(self, table_name: str, limit: Optional[int] = None) -> pd.DataFrame:
        """读取表数据"""
        if not self.connection:
            self.connect()
        
        try:
            query = f"SELECT * FROM {table_name}"
            if limit:
                query += f" LIMIT {limit}"
            
            df = pd.read_sql_query(query, self.connection)
            logger.info(f"成功读取表 {table_name}，共 {len(df)} 行数据")
            return df
        except Exception as e:
            logger.error(f"读取表 {table_name} 失败: {str(e)}")
            raise
    
    def read_query(self, query: str) -> pd.DataFrame:
        """执行自定义查询"""
        if not self.connection:
            self.connect()
        
        try:
            df = pd.read_sql_query(query, self.connection)
            logger.info(f"成功执行查询，共 {len(df)} 行数据")
            return df
        except Exception as e:
            logger.error(f"执行查询失败: {str(e)}")
            raise
    
    def get_table_schema(self, table_name: str) -> Dict[str, Any]:
        """获取表结构信息"""
        if not self.connection:
            self.connect()
        
        try:
            query = """
            SELECT column_name, data_type, is_nullable, column_default
            FROM information_schema.columns
            WHERE table_name = %s
            ORDER BY ordinal_position
            """
            
            cursor = self.connection.cursor()
            cursor.execute(query, (table_name,))
            columns = cursor.fetchall()
            
            schema = {
                'table_name': table_name,
                'columns': [
                    {
                        'name': col[0],
                        'type': col[1],
                        'nullable': col[2] == 'YES',
                        'default': col[3]
                    }
                    for col in columns
                ]
            }
            
            logger.info(f"成功获取表 {table_name} 的结构信息")
            return schema
        except Exception as e:
            logger.error(f"获取表结构失败: {str(e)}")
            raise
    
    def list_tables(self) -> List[str]:
        """获取数据库中所有表名"""
        if not self.connection:
            self.connect()
        
        try:
            query = """
            SELECT table_name
            FROM information_schema.tables
            WHERE table_schema = 'public'
            ORDER BY table_name
            """
            
            cursor = self.connection.cursor()
            cursor.execute(query)
            tables = [row[0] for row in cursor.fetchall()]
            
            logger.info(f"数据库中共有 {len(tables)} 个表")
            return tables
        except Exception as e:
            logger.error(f"获取表列表失败: {str(e)}")
            raise
    
    def __enter__(self):
        """上下文管理器入口"""
        self.connect()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器出口"""
        self.disconnect()

# 使用示例
if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(level=logging.INFO)
    
    # 使用数据读取器
    with DataReader() as reader:
        # 获取所有表
        tables = reader.list_tables()
        print(f"数据库中的表: {tables}")
        
        # 读取特定表的数据
        if tables:
            df = reader.read_table(tables[0], limit=10)
            print(f"表 {tables[0]} 的前10行数据:")
            print(df.head()) 